from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage, AIMessage

from common_config import CHAT_OLLAMA_MODEL

# 定义提示模板，要求模型直接输出答案
template = """
问题：{question}
答案：
/nothink
"""
prompt = ChatPromptTemplate.from_template(template)

# 初始化模型，base_url 默认为 http://localhost:11434，可不传
llm = ChatOllama(
    model=CHAT_OLLAMA_MODEL,
    temperature=0,
)  # model 为你已下载的模型名称


# 创建链
no_think = prompt | llm

# 调用模型生成回复（流式模式）
messages = {"question": "人工智能是什么？用少于100个字回答。"}

final_response = ""
# 流式输出
for chunk in no_think.stream(messages):
    print(chunk.content, end="", flush=True)
    final_response += chunk.content

print("\n\n回复总字符：" + str(len(final_response)))

# 定义提示模板
prompt = ChatPromptTemplate.from_template("你是一个专业厨师。请用简单步骤解释如何烹饪：{dish}。")

# 创建对话链：模板 -> 模型 -> 输出
cooker = prompt | no_think

# 调用链并生成回答
for chunk in cooker.stream({"dish": "红烧肉"}):
    print(chunk.content, end="", flush=True)


# 初始化对话历史
messages = [
    HumanMessage(content="你好！"),
    AIMessage(content="你好！有什么可以帮助你的？"),
    HumanMessage(content="推荐一本科幻小说。")
]

# 继续对话
for chunk in no_think.stream(messages):
    print(chunk.content, end="", flush=True)
